# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import functools
import logging
import os
from collections import namedtuple
from collections import deque

from coin.strategy.mm.dumper_base import (get_subscription_key,
                                          get_flow_subscriber_from_archive,
                                          get_flow_book_builder)
from coin.base.datetime_util import to_datetime, to_timestamp_int
from coin.base.book.types import TradeSide
from coin.exchange.base.order_gateway import OrderSide
from coin.strategy.mm.calculator import MovingSum
from coin.exchange.base.tick import FixedTickPrice


def run_from_archive(product_1,
                     product_2,
                     on_book_reset_callback,
                     from_ts,
                     to_ts=None,
                     worker_id='1',
                     machine=None,
                     root_dir=None):
  logging.info('Symbols: %s, %s' % (str(product_1), str(product_2)))
  flow_sub = get_flow_subscriber_from_archive([product_1, product_2],
                                              from_ts,
                                              machine=machine,
                                              root_dir=root_dir,
                                              worker_id=worker_id)

  for p in (product_1, product_2):
    sub_key = get_subscription_key(p)
    book_builder = get_flow_book_builder(sub_key, flow_sub, worker_id=worker_id)
    on_book_reset_callback(sub_key, book_builder)
  # Run
  flow_sub.start(until=to_ts)


class SimpleMatcher(object):
  _Fill = namedtuple('_Fill', ['price', 'qty', 'side'])
  _Order = namedtuple('_Order', ['price', 'qty', 'timestamp'])

  def __init__(self, side, execution_delay=0, post_only=False):
    self.side = side
    self.execution_delay = execution_delay
    self._post_only = post_only

    self.price = 0
    self.qty = 0
    self._order_queue = deque()
    self._last_ask0 = None
    self._last_bid0 = None

  def set_order(self, price, qty, timestamp):
    if ((self._order_queue and timestamp - self._order_queue[-1].timestamp < 2. * (10.**9))):
      return

    order = self._Order(price, qty, timestamp)
    self._order_queue.append(order)

  def cancel(self, timestamp):
    if self._order_queue and self._order_queue[-1].qty == 0:
      return

    order = self._Order(0, 0, timestamp)
    self._order_queue.append(order)

  def make_fill(self, fill_qty):
    assert fill_qty <= self.qty
    if fill_qty == 0:
      return None
    self.qty -= fill_qty
    return self._Fill(self.price, fill_qty, self.side)

  def _update_time(self, timestamp):
    updated = False
    while (self._order_queue
           and self._order_queue[0].timestamp + self.execution_delay <= timestamp):
      self.price = self._order_queue[0].price
      self.qty = self._order_queue[0].qty
      self._order_queue.popleft()
      updated = True

    if updated and self._post_only:
      last_midp = (self._last_ask0 + self._last_bid0) / 2.0
      if (((self.side == OrderSide.BUY and self.price >= last_midp)
           or (self.side == OrderSide.SELL and self.price <= last_midp))):
        self.qty = 0  # Auto-canceled

  def update_book(self, ask0, bid0, timestamp):
    self._update_time(timestamp)
    self._last_ask0 = ask0
    self._last_bid0 = bid0

    if self.qty == 0:
      return

    if self.side == OrderSide.BUY and self.price > ask0:
      return self.make_fill(self.qty)

    elif self.side == OrderSide.SELL and self.price < bid0:
      return self.make_fill(self.qty)

  def update_trade(self, trade_price, trade_qty, trade_side, timestamp):
    self._update_time(timestamp)
    fill_qty = min(trade_qty, self.qty)
    if fill_qty == 0:
      return

    if self.side == OrderSide.BUY and trade_side == TradeSide.TRADE_SELL_SIDE:
      if self.price > self._last_bid0:
        if self.price >= trade_price:
          return self.make_fill(fill_qty)
      else:
        if self.price > trade_price:
          return self.make_fill(fill_qty)

    elif self.side == OrderSide.SELL and trade_side == TradeSide.TRADE_BUY_SIDE:
      if self.price < self._last_ask0:
        if self.price <= trade_price:
          return self.make_fill(fill_qty)
      else:
        if self.price < trade_price:
          return self.make_fill(fill_qty)


class PositionCalculator(object):
  def __init__(self, fee):
    self.fee = fee

    self.pos = 0
    self.realized_pnl = 0
    self.unrealized_pnl = 0
    self.num_trade = 0
    self.num_buy = 0
    self.num_sell = 0
    self.qty_sum = 0

  def update_book(self, ask0, bid0):
    self._last_ask0 = ask0
    self._last_bid0 = bid0
    self._update_unrealized_pnl()

  def _update_unrealized_pnl(self):
    mid_price = (self._last_ask0 + self._last_bid0) / 2.
    self.unrealized_pnl = self.pos * mid_price

  def __str__(self):
    return (
        '[pnl:%.2f=%.2f+%.2f, pos:%.4f, num_trade:%d, qty_sum:%.4f]' %
        (self.pnl, self.realized_pnl, self.unrealized_pnl, self.pos, self.num_trade, self.qty_sum))

  @property
  def pnl(self):
    return self.realized_pnl + self.unrealized_pnl

  def update(self, price, qty, side):
    if side == OrderSide.BUY:
      self.pos += qty
      self.realized_pnl -= (price * qty) * (1. + self.fee)
      self.num_buy += 1
    elif side == OrderSide.SELL:
      self.pos -= qty
      self.realized_pnl += (price * qty) * (1. - self.fee)
      self.num_sell += 1
    else:
      raise ValueError(side)

    self.num_trade += 1
    self.qty_sum += qty
    self._update_unrealized_pnl()

  def update_by_fill_object(self, *fills):
    filled = []
    for fill in fills:
      if fill is not None:
        self.update(fill.price, fill.qty, fill.side)
        filled.append(fill)
    return filled


class SimExecutor(object):
  def __init__(self):
    self._products = {}
    self._exchange_ids = {}  # book_builder_name -> exchange_id
    self._bid_matchers = {}
    self._ask_matchers = {}
    self._pos = {}
    self._books = {}
    self._trade_qty_funcs = {}  # exchange_id -> trade -> qty

    self._book_callback = lambda book_1, book_2: None
    self._fill_callback = lambda exchange_id, price, qty, side, timestamp: None

  @property
  def initialized(self):
    return 1 in self._products and 2 in self._products

  def set_exchange(self,
                   exchange_id,
                   book_builder_name,
                   product,
                   fee,
                   execution_delay_sec=1.0,
                   post_only=False,
                   trade_qty_func=lambda trade: trade.qty):
    assert exchange_id in (1, 2), exchange_id
    self._products[exchange_id] = product
    self._exchange_ids[book_builder_name] = exchange_id
    self._bid_matchers[exchange_id] = SimpleMatcher(OrderSide.BUY,
                                                    execution_delay_sec * (10**9),
                                                    post_only)
    self._ask_matchers[exchange_id] = SimpleMatcher(OrderSide.SELL,
                                                    execution_delay_sec * (10**9),
                                                    post_only)
    self._trade_qty_funcs[exchange_id] = trade_qty_func
    self._pos[exchange_id] = PositionCalculator(fee)
    self._books[exchange_id] = None

  def on_book_reset(self, book_builder_name, book_builder):
    assert self.initialized
    exchange_id = self._exchange_ids.get(book_builder_name, None)
    assert exchange_id, book_builder_name
    book_builder.subscribe(self._products[exchange_id],
                           functools.partial(self.on_book, exchange_id))
    book_builder.subscribe_trade(self._products[exchange_id],
                                 functools.partial(self.on_trade, exchange_id))

  @property
  def pnl_sum(self):
    return sum([p.pnl for p in self._pos.values()])

  def on_book(self, exchange_id, book):
    self._books[exchange_id] = book

    try:
      ask0_p, bid0_p = book.ask0().price, book.bid0().price
    except IndexError:
      return
    self._pos[exchange_id].update_book(ask0_p, bid0_p)
    fill_bid = self._bid_matchers[exchange_id].update_book(ask0_p, bid0_p, book.timestamp)
    fill_ask = self._ask_matchers[exchange_id].update_book(ask0_p, bid0_p, book.timestamp)
    filled = self._pos[exchange_id].update_by_fill_object(fill_bid, fill_ask)

    for fill in filled:
      self._fill_callback(exchange_id, fill.price, fill.qty, fill.side, book.timestamp)
    if self._books[1] is not None and self._books[2] is not None:
      self._book_callback(self._books[1], self._books[2])

  def on_trade(self, exchange_id, trade):
    trade_qty = self._trade_qty_funcs[exchange_id](trade)

    bid_matcher = self._bid_matchers[exchange_id]
    ask_matcher = self._ask_matchers[exchange_id]
    fill_bid = bid_matcher.update_trade(trade.price, trade_qty, trade.side, trade.timestamp)
    fill_ask = ask_matcher.update_trade(trade.price, trade_qty, trade.side, trade.timestamp)
    filled = self._pos[exchange_id].update_by_fill_object(fill_bid, fill_ask)

    for fill in filled:
      self._fill_callback(exchange_id, fill.price, fill.qty, fill.side, trade.timestamp)

  def get_summary(self):
    return {
        'pnl': self.pnl_sum,
        'pos_1': self._pos[1].pos,
        'realized_pnl_1': self._pos[1].realized_pnl,
        'unrealized_pnl_1': self._pos[1].unrealized_pnl,
        'num_trade_1': self._pos[1].num_trade,
        'num_buy_1': self._pos[1].num_buy,
        'num_sell_1': self._pos[1].num_sell,
        'qty_sum_1': self._pos[1].qty_sum,
        'pos_2': self._pos[2].pos,
        'realized_pnl_2': self._pos[2].realized_pnl,
        'unrealized_pnl_2': self._pos[2].unrealized_pnl,
        'num_trade_2': self._pos[2].num_trade,
        'num_buy_2': self._pos[2].num_buy,
        'num_sell_2': self._pos[2].num_sell,
        'qty_sum_2': self._pos[2].qty_sum,
    }


class SimStrategy(object):
  def __init__(
      self,
      executor,
      basis_ma_window,  # Pricing
      reserve,  # Qty and execution
      imbalance_allowance,
      net_pos_allowance=10000,
      postfix=None,
      start_ts=None):
    self._basis_ma_window = basis_ma_window
    self._reserve = reserve

    # TODO(inkyu): Use inheritance
    self._book_askt_funcs = {}  # exchange_id -> book -> true_ask_price
    self._book_bidt_funcs = {}  # exchange_id -> book -> true_bid_price

    self._imbalance_allowance = imbalance_allowance
    self._net_pos_allowance = net_pos_allowance

    self._start_ts = to_timestamp_int(start_ts or 0)

    self._executor = None
    self._mid_basis_ma = MovingSum(self._basis_ma_window)

    self._edges = {}
    self._lot_sizes = {}
    self._ticks = {}

    self._fill_dump = None
    self._feature_dump = None

    self._bs_pos = 0.
    self._sb_pos = 0.
    self._postfix = postfix or ''
    self._register_executor(executor)

  def init_dump(self, dump_dir='.'):
    assert self.initialized

    param_str = self.get_filename()
    fill_filepath = os.path.join(dump_dir, 'fill_%s.csv' % param_str)
    self._fill_dump = open(fill_filepath, 'w')
    feature_filepath = os.path.join(dump_dir, 'feature_%s.csv' % param_str)
    self._feature_dump = open(feature_filepath, 'w')

    fill_header = ('time,time_hr,exchange,side,price,qty,'
                   'pos_1,pos_2,bs_pos,sb_pos,pnl_1,pnl_2,pnl\n')
    self._fill_dump.write(fill_header)

    self._last_feature_dump_ts = 0
    feature_header = ('time,time_hr,'
                      'ask0_1,bid0_1,askt_1,bidt_1,midt_1,'
                      'ask0_2,bid0_2,askt_2,bidt_2,midt_2,'
                      'mid_basis_ma,askp_1,bidp_1,askp_2,bidp_2\n')
    self._feature_dump.write(feature_header)

  def set_exchange(self,
                   exchange_id,
                   edge,
                   lot_size,
                   tick,
                   book_askt_func=lambda book: None,
                   book_bidt_func=lambda book: None):
    self._edges[exchange_id] = edge
    self._lot_sizes[exchange_id] = lot_size
    if isinstance(tick, float) or isinstance(tick, int):
      self._ticks[exchange_id] = FixedTickPrice(tick)
    else:
      self._ticks[exchange_id] = tick
    self._book_askt_funcs[exchange_id] = book_askt_func
    self._book_bidt_funcs[exchange_id] = book_bidt_func

  @property
  def initialized(self):
    return 1 in self._edges and 2 in self._edges

  def _register_executor(self, executor):
    self._executor = executor
    self._executor._book_callback = self.on_book
    self._executor._fill_callback = self.on_fill

  def get_filename(self):
    return ('ma%d_edge%gx%g_r%g_lot%gx%g_ia%g%s' %
            (self._basis_ma_window / (60. * 10**9),
             self._edges[1] * 10000,
             self._edges[2] * 10000,
             self._reserve,
             self._lot_sizes[1],
             self._lot_sizes[2],
             self._imbalance_allowance, ('_' + self._postfix) if self._postfix else ''))

  def __del__(self):
    if self._fill_dump:
      self._fill_dump.close()
    if self._feature_dump:
      self._feature_dump.close()

  def on_fill(self, exchange_id, fill_price, fill_qty, fill_side, timestamp):
    if exchange_id == 1:
      if fill_side == OrderSide.BUY:
        self._bs_pos += fill_qty
      elif fill_side == OrderSide.SELL:
        self._sb_pos -= fill_qty
    elif exchange_id == 2:
      if fill_side == OrderSide.BUY:
        self._sb_pos += fill_qty
      elif fill_side == OrderSide.SELL:
        self._bs_pos -= fill_qty
    else:
      raise ValueError(exchange_id)

    pos = self._executor._pos
    side_sign = 1 if fill_side == OrderSide.BUY else -1
    out = ('%d,"%s",%d,%d,%f,%f,%f,%f,%f,%f,%f,%f,%f\n' % (timestamp,
                                                           to_datetime(timestamp),
                                                           exchange_id,
                                                           side_sign,
                                                           fill_price,
                                                           fill_qty,
                                                           pos[1].pos,
                                                           pos[2].pos,
                                                           self._bs_pos,
                                                           self._sb_pos,
                                                           pos[1].pnl,
                                                           pos[2].pnl,
                                                           pos[1].pnl + pos[2].pnl))
    self._fill_dump.write(out)

  def on_book(self, book_1, book_2):
    assert self.initialized

    book_ts = max(book_1.timestamp, book_2.timestamp)
    min_ts = min(book_1.timestamp, book_2.timestamp)
    if book_ts - min_ts > (2 * 10**9):
      return

    # Exchange 1
    edge_1 = self._edges[1]
    tick_1 = self._ticks[1]
    ask0_1 = book_1.ask0().price
    bid0_1 = book_1.bid0().price
    askt_1 = self._book_askt_funcs[1](book_1)
    bidt_1 = self._book_bidt_funcs[1](book_1)
    midt_1 = (askt_1 + bidt_1) / 2.

    # Exchange 2
    edge_2 = self._edges[2]
    tick_2 = self._ticks[2]
    ask0_2 = book_2.ask0().price
    bid0_2 = book_2.bid0().price
    askt_2 = self._book_askt_funcs[2](book_2)
    bidt_2 = self._book_bidt_funcs[2](book_2)
    midt_2 = (askt_2 + bidt_2) / 2.

    mid_diff = midt_1 - midt_2
    self._mid_basis_ma.update(book_ts, mid_diff)

    if book_ts < self._start_ts:
      return

    if len(self._mid_basis_ma) < 2000:
      return

    midt_2_proj_1 = midt_2 + float(self._mid_basis_ma.average)
    midt_1_proj_2 = midt_1 - float(self._mid_basis_ma.average)

    ask0_minus_one_1 = tick_1.get_prev_price(ask0_1)
    bid0_plus_one_1 = tick_1.get_next_price(bid0_1)
    ask0_minus_one_2 = tick_2.get_prev_price(ask0_2)
    bid0_plus_one_2 = tick_2.get_next_price(bid0_2)

    askp_1 = max(tick_1.ceil(midt_2_proj_1 * (1. + edge_1 / 2.)), ask0_minus_one_1, bid0_plus_one_1)
    bidp_1 = min(tick_1.floor(midt_2_proj_1 * (1. - edge_1 / 2.)),
                 ask0_minus_one_1,
                 bid0_plus_one_1)
    askp_2 = max(tick_2.ceil(midt_1_proj_2 * (1. + edge_2 / 2.)), ask0_minus_one_2, bid0_plus_one_2)
    bidp_2 = min(tick_2.floor(midt_1_proj_2 * (1. - edge_2 / 2.)),
                 ask0_minus_one_2,
                 bid0_plus_one_2)

    if book_ts - self._last_feature_dump_ts >= 500 * (10**6):
      out = ('%d,"%s",%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f\n' %
             (book_ts,
              to_datetime(book_ts),
              ask0_1,
              bid0_1,
              askt_1,
              bidt_1,
              midt_1,
              ask0_2,
              bid0_2,
              askt_2,
              bidt_2,
              midt_2,
              float(self._mid_basis_ma.average),
              askp_1,
              bidp_1,
              askp_2,
              bidp_2))
      self._feature_dump.write(out)
      self._last_feature_dump_ts = book_ts

    self.post(2, askp_2, bidp_2, book_ts)
    self.post(1, askp_1, bidp_1, book_ts)

  def post(self, exchange_id, askp, bidp, timestamp):
    lot_size = self._lot_sizes[exchange_id]

    order_timestamp = timestamp + 10**6  # 1ms processing time
    no_buy = False
    net_pos = sum([pos.pos for pos in self._executor._pos.values()])

    if ((self._executor._pos[exchange_id].pos >= self._reserve)
        or (net_pos >= self._net_pos_allowance)
        or (exchange_id == 1 and self._bs_pos >= self._imbalance_allowance)
        or (exchange_id == 2 and self._sb_pos >= self._imbalance_allowance)):
      no_buy = True

    if no_buy:
      self._executor._bid_matchers[exchange_id].cancel(order_timestamp)
    else:
      self._executor._bid_matchers[exchange_id].set_order(bidp, lot_size, order_timestamp)

    no_sell = False
    if ((self._executor._pos[exchange_id].pos <= -self._reserve)
        or (net_pos <= -self._net_pos_allowance)
        or (exchange_id == 1 and self._sb_pos <= -self._imbalance_allowance)
        or (exchange_id == 2 and self._bs_pos <= -self._imbalance_allowance)):
      no_sell = True

    if no_sell:
      self._executor._ask_matchers[exchange_id].cancel(order_timestamp)
    else:
      self._executor._ask_matchers[exchange_id].set_order(askp, lot_size, order_timestamp)
